# -*- coding: utf-8 -*-
"""
-------------------------------------------------
Project Name: 分类项目，所有的内容整合在一个程序中
File Name: ui.py.py
Author: chenming
Create Date: 联系qq-3045834499，远程调试粉丝88元
Description：
-------------------------------------------------
"""
import sys
import shutil
import time
import os.path as osp
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from utils.model_operation import CLASSFICATION_MODEL
from utils.transforms import *

# 训练主要流程 训练的主流程由于没法实现进度条的同步，所以这个位置暂时关闭。
def train(train_loader, model, criterion, optimizer, learning_rate, epoch, total_epochs, device):
    # print("主函数搞定")
    metric_monitor = MetricMonitor()  # 设置指标监视器
    model.train()  # 模型设置为训练模型
    nBatch = len(train_loader)
    print(nBatch)
    stream = tqdm(train_loader)
    for i, (images, target) in enumerate(stream, start=1):  # 开始训练
        images = images.to(device, non_blocking=True)  # 加载数据
        target = target.to(device, non_blocking=True)  # 加载模型
        output = model(images)  # 数据送入模型进行前向传播
        loss = criterion(output, target.long())  # 计算损失
        f1_macro = calculate_f1_macro(output, target)  # 计算f1分数
        recall_macro = calculate_recall_macro(output, target)  # 计算recall分数
        acc = accuracy(output, target)  # 计算准确率分数
        metric_monitor.update('Loss', loss.item())  # 更新损失
        metric_monitor.update('F1', f1_macro)  # 更新f1
        metric_monitor.update('Recall', recall_macro)  # 更新recall
        metric_monitor.update('Accuracy', acc)  # 更新准确率
        optimizer.zero_grad()  # 清空学习率
        loss.backward()  # 损失反向传播
        optimizer.step()  # 更新优化器
        lr = adjust_learning_rate(optimizer, epoch, learning_rate, total_epochs, i, nBatch)  # 调整学习率
        # stream.set_description(  # 更新进度条
        #     "Epoch: {epoch}. Train.      {metric_monitor}".format(
        #         epoch=epoch,
        #         metric_monitor=metric_monitor)
        # )
    return metric_monitor.metrics['Accuracy']["avg"], metric_monitor.metrics['Loss']["avg"]  # 返回结果

# 验证流程
def validate(val_loader, model, criterion, device):
    metric_monitor = MetricMonitor()  # 验证流程
    model.eval()  # 模型设置为验证格式
    stream = tqdm(val_loader)  # 设置进度条
    with torch.no_grad():  # 开始推理
        for i, (images, target) in enumerate(stream, start=1):
            images = images.to(device, non_blocking=True)  # 读取图片
            target = target.to(device, non_blocking=True)  # 读取标签
            output = model(images)  # 前向传播
            loss = criterion(output, target.long())  # 计算损失
            f1_macro = calculate_f1_macro(output, target)  # 计算f1分数
            recall_macro = calculate_recall_macro(output, target)  # 计算recall分数
            acc = accuracy(output, target)  # 计算acc
            metric_monitor.update('Loss', loss.item())  # 后面基本都是更新进度条的操作
            metric_monitor.update('F1', f1_macro)
            metric_monitor.update("Recall", recall_macro)
            metric_monitor.update('Accuracy', acc)
            # stream.set_description(
            #     "Epoch: {epoch}. Validation. {metric_monitor}".format(
            #         epoch=epoch,
            #         metric_monitor=metric_monitor)
            # )
    return metric_monitor.metrics['Accuracy']["avg"], metric_monitor.metrics['Loss']["avg"]

# 训练进程
class Train_Worker(QObject):
    """在子线程中执行的工作对象"""
    progress = pyqtSignal(int)      # 发射进度 0~100
    message = pyqtSignal(str)       # 状态信息
    finished = pyqtSignal()         # 结束信号（正常完成或被取消）
    def __init__(self, model_name, model, epochs, train_loader,val_loader, device, learning_rate, model_save_path="", total=10, delay=0.03):
        super().__init__()
        print("线程初始化！！！！！！！！")
        self._total = total
        self._delay = delay
        self._stop = False
        self._paused = False
        # 模型的相关内容
        self.model_name = model_name
        self.model = model
        self.epochs = epochs
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        # 这这里需要将参数初始化好
        self.criterion = nn.CrossEntropyLoss().to(self.device)  # 设置损失函数
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=learning_rate)  # 设置优化器
        self.best_acc = 0.0
        self.model_save_path = model_save_path
        self.learning_rate = learning_rate


    @pyqtSlot()
    def run(self):
        # 将训练的整个过程集成到run里面
        # 整个for循环的过程植入到这个里面
        self.message.emit("模型训练中…")
        # 模型训练记录
        best_acc = 0
        accs = []
        losss = []
        val_accs = []
        val_losss = []
        nBatch = len(self.train_loader)
        # 模型重新和设备联动起来
        self.model = self.model.to(self.device)
        for epoch in range(1, self.epochs + 1):  # 开始训练

            # 真实的训练流程
            # acc, loss = train(self.train_loader, self.model, self.criterion, self.optimizer, self.learning_rate,
            #                   epoch, self.epochs, device=torch.device("cpu"))
            for i, (images, target) in enumerate(self.train_loader, start=1):  # 开始训练
                # 这里对应的是取消的操作
                if self._stop:
                    self.message.emit("已取消")
                    break
                # 暂停循环
                while self._paused and not self._stop:
                    time.sleep(0.05)

                images = images.to(self.device, non_blocking=True)  # 加载数据
                target = target.to(self.device, non_blocking=True)  # 加载模型
                output = self.model(images)  # 数据送入模型进行前向传播
                loss = self.criterion(output, target.long())  # 计算损失
                f1_macro = calculate_f1_macro(output, target)  # 计算f1分数
                recall_macro = calculate_recall_macro(output, target)  # 计算recall分数
                acc = accuracy(output, target)  # 计算准确率分数
                # metric_monitor.update('Loss', loss.item())  # 更新损失
                # metric_monitor.update('F1', f1_macro)  # 更新f1
                # metric_monitor.update('Recall', recall_macro)  # 更新recall
                # metric_monitor.update('Accuracy', acc)  # 更新准确率
                self.optimizer.zero_grad()  # 清空学习率
                loss.backward()  # 损失反向传播
                self.optimizer.step()  # 更新优化器
                self.learning_rate = adjust_learning_rate(self.optimizer, epoch, self.learning_rate, self.epochs, i, nBatch)  # 调整学习率
                accs.append(acc)
                losss.append(loss)
                # self.progress.emit(epoch-1 + i/nBatch)
                self.progress.emit(epoch)
            # 对模型状态进行判断
            # 这里对应的是取消的操作
            if self._stop:
                self.message.emit("已取消")
                break
            # 暂停循环
            while self._paused and not self._stop:
                time.sleep(0.05)
            # 每训练一轮之后要进行验证，如果不开启就不验证
            # 经过一轮的训练之后开始验证流程，查看当前模型的性能，另一方面为了后面在图形化界面上实时进行绘图
            val_acc, val_loss = validate(self.val_loader, self.model, self.criterion, self.device)
            val_accs.append(val_acc)
            val_losss.append(val_loss)
            if val_acc >= best_acc:
                # 保存的时候设置一个保存的间隔，或者就按照目前的情况，如果前面的比后面的效果好，就保存一下。
                # 按照间隔保存的话得不到最好的模型
                save_path = osp.join(self.model_save_path,
                                     f"{self.model_name}-{epoch}epochs_accuracy{val_acc:.5f}_weights.pth")
                torch.save(self.model.state_dict(), save_path)
                best_acc = val_acc

            time.sleep(0.05)
            self.progress.emit(epoch)
            time.sleep(self._delay)
        self.finished.emit()

    # 以下槽函数供主线程调用，修改运行状态
    @pyqtSlot()
    def pause(self):
        if not self._paused:
            self._paused = True
            self.message.emit("已暂停")

    @pyqtSlot()
    def resume(self):
        if self._paused:
            self._paused = False
            self.message.emit("继续执行…")

    @pyqtSlot()
    def stop(self):
        self._stop = True
        # 若在暂停中，也让其跳出暂停循环
        self._paused = False
        self.message.emit("正在取消…")

# 编写一个推理的接口
class Inference_Worker(QObject):
    """在子线程中执行的工作对象"""
    progress = pyqtSignal(int)      # 发射进度 0~100
    message = pyqtSignal(str)       # 状态信息
    finished = pyqtSignal()         # 结束信号（正常完成或被取消）
    def __init__(self, model_name, model, epochs, train_loader,val_loader, device, learning_rate, model_save_path="", total=10, delay=0.03):
        super().__init__()
        print("线程初始化！！！！！！！！")
        self._total = total
        self._delay = delay
        self._stop = False
        self._paused = False
        # 模型的相关内容
        self.model_name = model_name
        self.model = model
        self.epochs = epochs
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        # 这这里需要将参数初始化好
        self.criterion = nn.CrossEntropyLoss().to(self.device)  # 设置损失函数
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=learning_rate)  # 设置优化器
        self.best_acc = 0.0
        self.model_save_path = model_save_path
        self.learning_rate = learning_rate


    @pyqtSlot()
    def run(self):
        # 将训练的整个过程集成到run里面
        # 整个for循环的过程植入到这个里面
        self.message.emit("模型训练中…")
        # 模型训练记录
        best_acc = 0
        accs = []
        losss = []
        val_accs = []
        val_losss = []
        nBatch = len(self.train_loader)
        # 模型重新和设备联动起来
        self.model = self.model.to(self.device)
        for epoch in range(1, self.epochs + 1):  # 开始训练

            # 真实的训练流程
            # acc, loss = train(self.train_loader, self.model, self.criterion, self.optimizer, self.learning_rate,
            #                   epoch, self.epochs, device=torch.device("cpu"))
            for i, (images, target) in enumerate(self.train_loader, start=1):  # 开始训练
                # 这里对应的是取消的操作
                if self._stop:
                    self.message.emit("已取消")
                    break
                # 暂停循环
                while self._paused and not self._stop:
                    time.sleep(0.05)

                images = images.to(self.device, non_blocking=True)  # 加载数据
                target = target.to(self.device, non_blocking=True)  # 加载模型
                output = self.model(images)  # 数据送入模型进行前向传播
                loss = self.criterion(output, target.long())  # 计算损失
                f1_macro = calculate_f1_macro(output, target)  # 计算f1分数
                recall_macro = calculate_recall_macro(output, target)  # 计算recall分数
                acc = accuracy(output, target)  # 计算准确率分数
                # metric_monitor.update('Loss', loss.item())  # 更新损失
                # metric_monitor.update('F1', f1_macro)  # 更新f1
                # metric_monitor.update('Recall', recall_macro)  # 更新recall
                # metric_monitor.update('Accuracy', acc)  # 更新准确率
                self.optimizer.zero_grad()  # 清空学习率
                loss.backward()  # 损失反向传播
                self.optimizer.step()  # 更新优化器
                self.learning_rate = adjust_learning_rate(self.optimizer, epoch, self.learning_rate, self.epochs, i, nBatch)  # 调整学习率
                accs.append(acc)
                losss.append(loss)
                # self.progress.emit(epoch-1 + i/nBatch)
                self.progress.emit(epoch)
            # 对模型状态进行判断
            # 这里对应的是取消的操作
            if self._stop:
                self.message.emit("已取消")
                break
            # 暂停循环
            while self._paused and not self._stop:
                time.sleep(0.05)
            # 每训练一轮之后要进行验证，如果不开启就不验证
            # 经过一轮的训练之后开始验证流程，查看当前模型的性能，另一方面为了后面在图形化界面上实时进行绘图
            val_acc, val_loss = validate(self.val_loader, self.model, self.criterion, self.device)
            val_accs.append(val_acc)
            val_losss.append(val_loss)
            if val_acc >= best_acc:
                # 保存的时候设置一个保存的间隔，或者就按照目前的情况，如果前面的比后面的效果好，就保存一下。
                # 按照间隔保存的话得不到最好的模型
                save_path = osp.join(self.model_save_path,
                                     f"{self.model_name}-{epoch}epochs_accuracy{val_acc:.5f}_weights.pth")
                torch.save(self.model.state_dict(), save_path)
                best_acc = val_acc

            time.sleep(0.05)
            self.progress.emit(epoch)
            time.sleep(self._delay)
        self.finished.emit()

    # 以下槽函数供主线程调用，修改运行状态
    @pyqtSlot()
    def pause(self):
        if not self._paused:
            self._paused = True
            self.message.emit("已暂停")

    @pyqtSlot()
    def resume(self):
        if self._paused:
            self._paused = False
            self.message.emit("继续执行…")

    @pyqtSlot()
    def stop(self):
        self._stop = True
        # 若在暂停中，也让其跳出暂停循环
        self._paused = False
        self.message.emit("正在取消…")

class MainWindow(QTabWidget):
    # 基本配置不动，然后只动第三个界面
    def __init__(self):
        super().__init__()
        # 初始化对象
        self.class_model_meta_info = CLASSFICATION_MODEL()
        # 从本地加载配置信息
        self.class_model_meta_info.update_model_info_from_local()
        # 系统配置信息另外做一个界面来进行保存
        # 下方是系统配置信息区
        self.setWindowTitle('分类系统')
        self.resize(1200, 800)
        self.setWindowIcon(QIcon("images/UI/lufei.png"))
        self.output_size = 480
        self.img2predict = ""
        self.origin_shape = ()
        # 1. 加载配置
        # 2. 根据配置文件加载对应的模型
        # 初始化图形化界面
        self.initUI()

        # 更新界面信息
        self.refresh_widget_config_info()

    '''
    ***界面初始化***
    '''
    def initUI(self):
        # 增加一个模型训练的子界面，可以设置参数，训练以及进度条
        ################################
        # 模型训练子页面
        ################################
        self.train_widget = QWidget()
        self.train_layout = QVBoxLayout()
        # 参数调节界面位置
        config_widget = QWidget()
        config_layout = QGridLayout()
        self.data_folder_label = QPushButton("选择数据集文件夹")
        self.data_folder_edit = QLineEdit("")
        self.data_folder_edit.setReadOnly(True)
        self.data_folder_edit.setPlaceholderText("未选择")
        self.data_folder_label.clicked.connect(self.choose_data_folder)

        # 根据数据集的路径信息获取类别数量和类别名称
        self.classes_names_label =QLabel("数据集类别")
        self.classes_names_edit = QComboBox()

        self.num_classes_label = QLabel("数据集类别数目")
        self.num_classes_edit = QLineEdit("")
        self.num_classes_label.setEnabled(False)

        self.image_size_label = QLabel("图像大小")
        self.image_size_edit = QLineEdit()

        self.batch_size_label = QLabel("批次大小")
        self.batch_size_edit = QLineEdit()

        self.model_name_label = QLabel("模型名称")
        self.model_name_eidt =QLineEdit()

        self.device_label = QLabel("所使用的设备")
        self.device_eidt = QLineEdit()

        self.learning_rate_label = QLabel("学习率")
        self.learning_rate_eidt = QLineEdit()

        self.epochs_label = QLabel("模型训练轮数")
        self.epochs_eidt = QLineEdit()

        self.model_save_path_label = QLabel("模型保存路径")
        self.model_save_path_eidt = QLineEdit()

        self.weight_decay_label = QLabel("权重衰减")
        self.weight_decay_eidt = QLineEdit()

        self.pretrained_label = QLabel("是否使用预训练")
        self.pretrained_eidt = QComboBox()
        self.pretrained_eidt.addItems(["False", "True"])
        
        # 配置信息添加到网格布局
        config_layout.addWidget(self.data_folder_label, 0, 0)
        config_layout.addWidget(self.data_folder_edit, 0, 1)
        config_layout.addWidget(self.classes_names_label)
        config_layout.addWidget(self.classes_names_edit)
        config_layout.addWidget(self.num_classes_label)
        config_layout.addWidget(self.num_classes_edit)
        config_layout.addWidget(self.image_size_label)
        config_layout.addWidget(self.image_size_edit)
        config_layout.addWidget(self.batch_size_label)
        config_layout.addWidget(self.batch_size_edit)
        config_layout.addWidget(self.model_name_label)
        config_layout.addWidget(self.model_name_eidt)
        config_layout.addWidget(self.device_label)
        config_layout.addWidget(self.device_eidt)
        config_layout.addWidget(self.learning_rate_label)
        config_layout.addWidget(self.learning_rate_eidt)
        config_layout.addWidget(self.epochs_label)
        config_layout.addWidget(self.epochs_eidt)
        config_layout.addWidget(self.model_save_path_label)
        config_layout.addWidget(self.model_save_path_eidt)
        config_layout.addWidget(self.weight_decay_label)
        config_layout.addWidget(self.weight_decay_eidt)
        config_layout.addWidget(self.pretrained_label)
        config_layout.addWidget(self.pretrained_eidt)
        # 网格组件上设置对应的网格布局
        config_widget.setLayout(config_layout)
        self.bar = QProgressBar()
        self.bar.setRange(0, 10)
        self.bar.setValue(0)
        self.train_label = QLabel("训练准备就绪")
        self.train_label.setAlignment(Qt.AlignCenter)
        self.btn_save_config = QPushButton("保存配置")
        self.btn_start = QPushButton("开始")
        self.btn_cancel = QPushButton("结束任务")
        # 初始按钮状态
        self.btn_cancel.setEnabled(False)
        h = QHBoxLayout()
        h.addWidget(self.btn_start)
        h.addWidget(self.btn_cancel)

        # lay = QVBoxLayout(self)
        self.train_layout.addWidget(config_widget)
        self.train_layout.addWidget(self.btn_save_config)
        self.train_layout.addWidget(self.bar)
        self.train_layout.addWidget(self.train_label)
        self.train_layout.addLayout(h)
        # 连接按钮信号
        self.btn_start.clicked.connect(self.train_start_task)
        self.btn_cancel.clicked.connect(self.train_cancel_task)
        self.btn_save_config.clicked.connect(self.save_config_info)
        # 线程/worker 占位
        self.thread = None
        self.worker = None
        # 设置界面之后添加到过程中
        self.train_widget.setLayout(self.train_layout)

        ######################## 推理界面 ##############################
        inference_widget = QWidget()
        inference_widget_layout = QHBoxLayout()

        # 左侧图像显示区域
        show_widget = QWidget()
        show_layout = QVBoxLayout()
        show_label = QLabel("推理结果显示区域")
        # 右侧参数区域和按钮区域
        self.result_image = QLabel("推理结果")
        show_layout.addWidget(show_label)
        show_layout.addWidget(self.result_image)

        inference_config_widget = QWidget()
        inference_config_widget_layout = QVBoxLayout()
        inference_config_title = QLabel("推理参数")

        # 推理参数部分使用网格布局，分类任务主要是选择对应的模型，以及图像推理的尺寸大小
        inference_params_widget = QWidget()
        inference_params_layout = QGridLayout()
        # 模型路径
        inference_model_label = QLabel("模型路径")
        self.inference_model_edit = QLineEdit("")
        # 图像推理设备
        inference_device_label = QLabel("所使用的设备")
        self.inference_device_edit = QLineEdit("")
        # 单图推理
        self.inference_image_btn = QPushButton("单图推理")
        self.inference_image_path = QLineEdit("")
        self.inference_image_path.setReadOnly(True)
        self.inference_image_path.setPlaceholderText("未选择")
        # 文件夹推理
        self.inference_folder_btn = QPushButton("文件夹推理")
        self.inference_folder_path = QLineEdit("")
        self.inference_folder_path.setReadOnly(True)
        self.inference_folder_path.setPlaceholderText("未选择")

        # 视频推理
        self.inference_video_btn = QPushButton("视频推理")
        self.inference_video_path = QLineEdit("")
        self.inference_video_path.setReadOnly(True)
        self.inference_video_path.setPlaceholderText("未选择")

        # 摄像头推理
        self.inference_webcam_btn = QPushButton("摄像头推理")
        self.inference_webcam_path = QLineEdit("")
        # self.inference_image_path.setReadOnly(True)
        self.inference_webcam_path.setPlaceholderText("未选择")

        inference_params_layout.addWidget(inference_model_label, 0, 0)
        inference_params_layout.addWidget(self.inference_model_edit, 0, 1)
        inference_params_layout.addWidget( inference_device_label)
        inference_params_layout.addWidget(self.inference_device_edit )
        inference_params_layout.addWidget(self.inference_image_path)
        inference_params_layout.addWidget(self.inference_image_btn)
        inference_params_layout.addWidget( self.inference_folder_path)
        inference_params_layout.addWidget(self.inference_folder_btn)
        inference_params_layout.addWidget(self.inference_video_path)
        inference_params_layout.addWidget(self.inference_video_btn)
        inference_params_layout.addWidget(self.inference_webcam_path)
        inference_params_layout.addWidget(self.inference_webcam_btn)
        inference_params_widget.setLayout(inference_params_layout)
        inference_config_widget_layout.addWidget(inference_config_title)
        inference_config_widget_layout.addWidget(inference_params_widget)
        inference_config_widget.setLayout(inference_config_widget_layout)

        inference_widget_layout.addWidget(show_widget)
        inference_widget_layout.addWidget(inference_config_widget)
        inference_widget.setLayout(inference_widget_layout)


        # inference_config_widget = QWidget()
        # inference_config_widget_layout = QVBoxLayout()
        # inference_config_title = QLabel("推理参数")
        # font_title = QFont('楷体', 16)
        # font_main = QFont('楷体', 14)
        # # 图片识别界面, 两个按钮，上传图片和显示结果
        # img_detection_widget = QWidget()
        # img_detection_layout = QVBoxLayout()
        # img_detection_title = QLabel("图片识别功能")
        # img_detection_title.setFont(font_title)
        # mid_img_widget = QWidget()
        # mid_img_layout = QHBoxLayout()
        # self.left_img = QLabel()
        # self.right_img = QLabel()
        # self.left_img.setPixmap(QPixmap("images/UI/up.jpeg"))
        # self.right_img.setPixmap(QPixmap("images/UI/right.jpeg"))
        # self.left_img.setAlignment(Qt.AlignCenter)
        # self.right_img.setAlignment(Qt.AlignCenter)
        # mid_img_layout.addWidget(self.left_img)
        # mid_img_layout.addStretch(0)
        # # mid_img_layout.addWidget(self.right_img)
        # mid_img_widget.setLayout(mid_img_layout)
        # up_img_button = QPushButton("上传图片")
        # det_img_button = QPushButton("识别")
        # up_img_button.clicked.connect(self.upload_img)
        # det_img_button.clicked.connect(self.detect_img)
        # up_img_button.setFont(font_main)
        # det_img_button.setFont(font_main)
        # self.rrr = QLabel("等待识别")
        # self.rrr.setFont(font_main)
        # up_img_button.setStyleSheet("QPushButton{color:white}"
        #                             "QPushButton:hover{background-color: rgb(2,110,180);}"
        #                             "QPushButton{background-color:rgb(48,124,208)}"
        #                             "QPushButton{border:2px}"
        #                             "QPushButton{border-radius:5px}"
        #                             "QPushButton{padding:5px 5px}"
        #                             "QPushButton{margin:5px 5px}")
        # det_img_button.setStyleSheet("QPushButton{color:white}"
        #                              "QPushButton:hover{background-color: rgb(2,110,180);}"
        #                              "QPushButton{background-color:rgb(48,124,208)}"
        #                              "QPushButton{border:2px}"
        #                              "QPushButton{border-radius:5px}"
        #                              "QPushButton{padding:5px 5px}"
        #                              "QPushButton{margin:5px 5px}")
        # img_detection_layout.addWidget(img_detection_title, alignment=Qt.AlignCenter)
        # img_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
        # img_detection_layout.addWidget(self.rrr)
        # # img_detection_layout.addWidget(self.c1)
        # img_detection_layout.addWidget(up_img_button)
        # img_detection_layout.addWidget(det_img_button)
        # img_detection_widget.setLayout(img_detection_layout)
        # 设置视频识别的位置可以对视频进行读取
        '''
        *** 关于界面 ***
        '''
        about_widget = QWidget()
        about_layout = QVBoxLayout()
        about_title = QLabel('欢迎使用基于Pytorch的分类系统')  # todo 修改欢迎词语
        about_title.setFont(QFont('楷体', 18))
        about_title.setAlignment(Qt.AlignCenter)
        about_img = QLabel()
        about_img.setPixmap(QPixmap('images/UI/qq.png'))
        about_img.setAlignment(Qt.AlignCenter)

        # label4.setText("<a href='https://oi.wiki/wiki/学习率的调整'>如何调整学习率</a>")
        label_super = QLabel()  # todo 更换作者信息
        label_super.setText("<a href='https://blog.csdn.net/ECHOSON'>或者你可以在这里找到我-->肆十二</a>")
        label_super.setFont(QFont('楷体', 16))
        label_super.setOpenExternalLinks(True)
        # label_super.setOpenExternalLinks(True)
        label_super.setAlignment(Qt.AlignRight)
        about_layout.addWidget(about_title)
        about_layout.addStretch()
        about_layout.addWidget(about_img)
        about_layout.addStretch()
        about_layout.addWidget(label_super)
        about_widget.setLayout(about_layout)

        # self.left_img.setAlignment(Qt.AlignCenter)
        self.addTab(self.train_widget, '模型训练')
        self.addTab(inference_widget, '模型推理')
        self.addTab(about_widget, '系统配置')
        self.setTabIcon(0, QIcon('images/UI/lufei.png'))
        self.setTabIcon(1, QIcon('images/UI/lufei.png'))

    def choose_data_folder(self):
        """
        选择数据集文件夹
        :return:
        """
        
        # 起始目录设为用户主目录
        start_dir = os.path.expanduser("~")
        # 仅显示文件夹，避免解析符号链接
        options = QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks
        folder = QFileDialog.getExistingDirectory(
            self, "选择文件夹", start_dir, options=options
        )
        if folder:  # 用户点了“确定”
            if 'train' in os.listdir(folder):
                # 更新数据集的路径，参数更新和界面更新
                self.data_folder_edit.setText(folder)
                # 更新对象里面的信息
                self.save_config_info()
                # 更新界面组件上的信息
                self.refresh_widget_config_info()
            else:
                QMessageBox.information(self, "错误", "您选择的文件夹下需要包含train目录")

    # 选择图片进行推理

    # 选择文件夹进行推理

    # 选在视频进行推理

    # 输入摄像头进行推理

    def refresh_widget_config_info(self):
        """
        刷新界面的配置信息
        :return:
        """
        self.data_folder_edit.setText(str(self.class_model_meta_info.data_folder))
        # 根据具体的类别名称去除原先的内容并添加新的内容。
        self.classes_names_edit.clear()
        self.classes_names_edit.addItems(self.class_model_meta_info.class_names)
        self.num_classes_edit.setText(str(self.class_model_meta_info.num_classes))
        self.image_size_edit.setText(str(self.class_model_meta_info.image_size))
        self.batch_size_edit.setText(str(self.class_model_meta_info.batch_size))
        self.model_name_eidt.setText(str(self.class_model_meta_info.model_name))
        self.device_eidt.setText(str(self.class_model_meta_info.device))
        self.learning_rate_eidt.setText(str(self.class_model_meta_info.learning_rate))
        self.epochs_eidt.setText(str(self.class_model_meta_info.epochs))
        self.model_save_path_eidt.setText(str(self.class_model_meta_info.model_save_path))
        self.weight_decay_eidt.setText(str(self.class_model_meta_info.weight_decay))
        # self.pretrained_eidt.setText(str(self.class_model_meta_info.pretrained))
        self.pretrained_eidt.setCurrentIndex(int(self.class_model_meta_info.pretrained))

    def save_config_info(self):
        # 直接保存即可，读取网格布局中组件的信息
        # 熟悉对象里面的信息
        self.class_model_meta_info.data_folder = self.data_folder_edit.text()
        # 其中数据集的类别和名称
        self.class_model_meta_info.class_names = os.listdir(os.path.join(self.class_model_meta_info.data_folder, "train"))
        self.class_model_meta_info.num_classes = len(self.class_model_meta_info.class_names)
        # 类别和应该使用下拉框表示出来
        # 如果数据集路径切换了，这些信息应该是联动的，直接根据数据集路径来更新信息即可。
        self.class_model_meta_info.image_size = int(self.image_size_edit.text())
        self.class_model_meta_info.batch_size = int(self.batch_size_edit.text())
        self.class_model_meta_info.model_name = self.model_name_eidt.text()
        self.class_model_meta_info.device =self.device_eidt.text()
        self.class_model_meta_info.learning_rate =float(self.learning_rate_eidt.text())
        self.class_model_meta_info.epochs =int(self.epochs_eidt.text())
        self.class_model_meta_info.model_save_path =self.model_save_path_eidt.text()
        self.class_model_meta_info.weight_decay =float(self.weight_decay_eidt.text())
        self.class_model_meta_info.pretrained = self.pretrained_eidt.currentText()
        if str(self.class_model_meta_info.pretrained).lower() == "false":
            self.class_model_meta_info.pretrained = False
        else:
            self.class_model_meta_info.pretrained= True
        # 将配置信息保存到本地
        self.class_model_meta_info.save_local_config()
        QMessageBox.information(self, "保存成功", "配置保存成功!")

    def train_start_task(self):
        if self.thread is not None:
            return  # 已在运行

        # 创建线程与工作者
        self.thread = QThread()
        # 获取参数并且将参数送入实例化之后执行
        # 先完成参数获取的任务，然后完成训练的任务
#
#         self.worker = Train_Worker(total=100, delay=0.03)
        train_loader, val_loader, test_loader = self.class_model_meta_info.data_init()
        print("数据加载完毕！！！！！！！！！！！！！")
        # 模型初始化完成
        self.worker = Train_Worker(self.class_model_meta_info.model_name, self.class_model_meta_info.model, self.class_model_meta_info.epochs, train_loader, val_loader, self.class_model_meta_info.device, self.class_model_meta_info.learning_rate, model_save_path="", total=10, delay=0.03)
        self.worker.moveToThread(self.thread)

        # 线程启动后调用 worker.run
        self.thread.started.connect(self.worker.run)

        # worker 信号回到 UI
        self.worker.progress.connect(self.bar.setValue)
        self.worker.message.connect(self.train_label.setText)
        self.worker.finished.connect(self.on_finished)

        # 生命周期管理：结束时退出线程
        self.worker.finished.connect(self.thread.quit)

        # 清理：线程结束后删除对象
        self.thread.finished.connect(self.thread.deleteLater)
        self.thread.finished.connect(self._clear_thread_refs)

        # UI 状态
        self.bar.setValue(0)
        self.train_label.setText("准备开始…")
        self.btn_start.setEnabled(False)
        # self.btn_pause.setEnabled(True)
        # self.btn_resume.setEnabled(False)
        self.btn_cancel.setEnabled(True)
        self.thread.start()

    def train_cancel_task(self):
        if self.worker:
            self.worker.stop()
            # 等待 on_finished 做 UI 恢复

    # def pause_task(self):
    #     if self.worker:
    #         self.worker.pause()
    #         self.btn_pause.setEnabled(False)
    #         self.btn_resume.setEnabled(True)

    # def resume_task(self):
    #     if self.worker:
    #         self.worker.resume()
    #         self.btn_pause.setEnabled(True)
    #         self.btn_resume.setEnabled(False)

    @pyqtSlot()
    def on_finished(self):
        self.train_label.setText("任务结束")
        self.btn_start.setEnabled(True)
        # self.btn_pause.setEnabled(False)
        # self.btn_resume.setEnabled(False)
        self.btn_cancel.setEnabled(False)

    def _clear_thread_refs(self):
        # 断开与删除 worker 引用（让 Qt 负责回收）
        self.worker.deleteLater()
        self.worker = None
        self.thread = None


    def upload_img(self):
        # 选择录像文件进行读取
        fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.jpg *.png *.tif *.jpeg')
        if fileName:
            suffix = fileName.split(".")[-1]
            save_path = osp.join("resources/images/tmp", "tmp_upload." + suffix)
            shutil.copy(fileName, save_path)
            # 应该调整一下图片的大小，然后统一防在一起
            im0 = cv2.imread(save_path)
            resize_scale = self.output_size / im0.shape[0]
            im0 = cv2.resize(im0, (0, 0), fx=resize_scale, fy=resize_scale)
            cv2.imwrite("resources/images/tmp/upload_show_result.jpg", im0)
            # self.right_img.setPixmap(QPixmap("images/tmp/single_result.jpg"))
            self.img2predict = fileName
            self.origin_shape = (im0.shape[1], im0.shape[0])
            self.left_img.setPixmap(QPixmap("images/tmp/upload_show_result.jpg"))
            # todo 上传图片之后右侧的图片重置，
            self.right_img.setPixmap(QPixmap("images/UI/right.jpeg"))
            self.rrr.setText("等待识别")


    '''
    ***检测图片***
    '''
    # 写一个通用的内容，在主界面上，包含一些对花卉的介绍。
    def detect_img(self):
        # model = self.model
        # output_size = self.output_size
        source = self.img2predict  # file/dir/URL/glob, 0 for webcam
        img = Image.open(source)
        img = img.convert('RGB')
        img = self.valid_transforms(img)
        img = img.unsqueeze(0)
        img = img.to(self.device)
        output = self.model(img)
        label_id = torch.argmax(output).item()
        predict_name = self.classes_names[label_id]
        self.rrr.setText("当前识别结果为：{}".format(predict_name))

    # 关闭事件 询问用户是否退出
    def closeEvent(self, event):
        reply = QMessageBox.question(self,
                                     '退出',
                                     "是否要退出程序？",
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.close()
            event.accept()
        else:
            event.ignore()

if __name__ == "__main__":
    # test_model = CLASSFICATION_MODEL()
    # test_model.update_model_info()
    # test_model.save_local_config()
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())
